import datetime
import asyncio
import base64
import struct
import sys
import urllib.parse
import os

class WebsocketSession:
    def __init__(self):
        self.reader = None
        self.writer = None
        self._id = "12345"  # 简单设置一个会话 ID
        self._is_closed = False

    @property
    def id(self) -> str:
        return self._id

    @property
    def is_closed(self) -> bool:
        return self._is_closed

    async def connect(self, endpoint):
        print(f'[system] Connecting to {endpoint}')
        parsed = urllib.parse.urlparse(endpoint)
        scheme = parsed.scheme
        host = parsed.hostname
        port = parsed.port
        path = parsed.path if parsed.path else "/"

        if scheme == "ws":
            if not port:
                port = 80
        elif scheme == "wss":
            if not port:
                port = 443
            raise NotImplementedError("wss (WebSocket Secure) is not implemented in this simple version.")
        else:
            raise ValueError(f"Unsupported scheme: {scheme}")

        self.reader, self.writer = await asyncio.open_connection(host, port)

        # WebSocket 握手
        key = base64.b64encode(os.urandom(16)).decode()
        headers = [
            f"GET {path} HTTP/1.1",
            f"Host: {host}:{port}",
            "Upgrade: websocket",
            "Connection: Upgrade",
            f"Sec-WebSocket-Key: {key}",
            "Sec-WebSocket-Version: 13",
            "\r\n"
        ]
        request = "\r\n".join(headers)
        self.writer.write(request.encode())
        await self.writer.drain()

        response = await self.reader.read(1024)
        response = response.decode()
        if "101 Switching Protocols" not in response:
            for line in response.split("\r\n"):
                print(f'>>> {line}')
            raise Exception("WebSocket handshake failed")
        print('>>> WebSocket handshake successful')

    def send(self, message):
        if isinstance(message, str):
            asyncio.create_task(self.send_text(message))
        else:
            asyncio.create_task(self.send_binary(message))

    async def send_text(self, message):
        await self._send_frame(0x01, message.encode())

    async def send_binary(self, binary):
        await self._send_frame(0x02, binary)

    def send_pong(self, message=b''):
        asyncio.create_task(self._send_frame(0x0A, message))

    def send_ping(self, message=b''):
        asyncio.create_task(self._send_frame(0x09, message))

    async def close(self, reason=""):
        if not self._is_closed:
            try:
                await self._send_frame(0x08, reason.encode())
                self._is_closed = True
                self.writer.close()
                await self.writer.wait_closed()
            except Exception as e:
                print(f"Error closing connection: {e}")

    async def _send_frame(self, opcode, payload):
        frame = bytearray()
        frame.append(0x80 | opcode)

        payload_length = len(payload)
        mask_bit = 0x80  # 表示使用掩码
        if payload_length < 126:
            frame.append(mask_bit | payload_length)
        elif payload_length < 65536:
            frame.append(mask_bit | 126)
            frame.extend(struct.pack('>H', payload_length))
        else:
            frame.append(mask_bit | 127)
            frame.extend(struct.pack('>Q', payload_length))

        # 生成 32 位掩码密钥
        masking_key = os.urandom(4)
        frame.extend(masking_key)

        # 对负载数据进行掩码处理
        masked_payload = bytes([payload[i] ^ masking_key[i % 4] for i in range(len(payload))])
        frame.extend(masked_payload)

        self.writer.write(frame)
        await self.writer.drain()

    async def receive(self):
        try:
            header = await self.reader.read(2)
            if not header:  # 如果没有收到任何数据，则认为连接已关闭
                return None
            fin = (header[0] & 0x80) == 0x80
            opcode = header[0] & 0x0F
            mask = (header[1] & 0x80) == 0x80
            payload_length = header[1] & 0x7F

            if payload_length == 126:
                payload_length = struct.unpack('>H', await self.reader.read(2))[0]
            elif payload_length == 127:
                payload_length = struct.unpack('>Q', await self.reader.read(8))[0]

            if mask:
                masking_key = await self.reader.read(4)

            payload = await self.reader.read(payload_length)
            if mask:
                payload = bytes([payload[i] ^ masking_key[i % 4] for i in range(len(payload))])

            if opcode == 0x08:  # Close frame
                self._is_closed = True
                self.writer.close()
                await self.writer.wait_closed()
            return payload
        except asyncio.TimeoutError:
            print("socket timeout")
            return None

async def task_ws_reader(ws_client: WebsocketSession):
    while not ws_client.is_closed:
        response = await ws_client.receive()
        if not response:
            print('websocket closed')
            break
        try:
            current_time = datetime.datetime.now()
            now = current_time.strftime("%Y-%m-%d %H:%M:%S")
            print(f"\r(server) {now}: {response.decode()}")
            print('[self] > ', end='', flush=True)
        except UnicodeDecodeError:
            print(f"Received binary data: {response}")

async def task_stdin_reader(ws_client: WebsocketSession):
    while not ws_client.is_closed:
        try:
            message = await asyncio.get_event_loop().run_in_executor(None, input, "[self] > ")
            ws_client.send(message)
        except EOFError:
            # pressed Ctrl+C
            print("bye!")
            await ws_client.close("User terminated input")
            break

async def main():
    endpoint = "ws://localhost:9090/ws/echo"
    if len(sys.argv) > 1:
        endpoint = sys.argv[1]
    client = WebsocketSession()
    try:
        await client.connect(endpoint)
        # begin forwarding between stdin/out <-> websocket
        await asyncio.gather(
            task_ws_reader(client),
            task_stdin_reader(client)
        )
    except asyncio.exceptions.CancelledError:
        print("User interrupted with Ctrl+C. Closing connection...")
        await client.close("User interrupted with Ctrl+C")
    except Exception as e:
        await client.close("unexpected error occurred")
        print(f"unexpected error occurred: {e}")
        raise

if __name__ == "__main__":
    asyncio.run(main())
